import pickle
import tqdm
from collections import Counter


class TorchVocab(object):
    """
        该类主要是定义了一个词典对象，包含如下三个属性：

        freqs:是一个collections.Counter对象,能够记录数据集中不同token所出现的次数
        stoi:是一个collections.defaultdict对象,将数据集中的token映射到不同的index
        itos:是一个列表,保存了从index到token的映射信息
    """

    def __init__(self, counter, max_size=None, min_freq=1, specials=['<pad>', '<oov>'],
                 vectors=None, unk_init=None, vectors_cache=None):
        """从 collections.Counter 创建一个 Vocab 对象。
        参数:
            counter: 一个 collections.Counter 对象，保存在数据中找到的每个值的频率。
            max_size: 词汇表的最大大小，或无最大限制时为 None。默认: None。
            min_freq: 包含标记在词汇表中的最小频率。小于 1 的值将被设置为 1。默认: 1。
            specials: 特殊标记的列表（例如，填充或结束标记），这些标记将被添加到词汇表中，除了一个 <unk> 标记。默认: ['<pad>']。
            vectors: 可以是可用的预训练向量之一或自定义预训练向量（参见 Vocab.load_vectors）；也可以是上述向量的列表。
            unk_init (回调): 默认情况下，将未出现在词汇表中的单词向量初始化为零向量；可以是任何接受一个 Tensor 并返回相同大小的 Tensor 的函数。默认: torch.Tensor.zero_
            vectors_cache: 预训练向量的缓存目录。默认: '.vector_cache'
        """
        self.freqs = counter
        counter = counter.copy()
        min_freq = max(min_freq, 1)

        self.itos = list(specials)
        # 构建按频率排序的词汇表时不计算特殊标记的频率
        for tok in specials:
            del counter[tok]

        max_size = None if max_size is None else max_size + len(self.itos)
        # 按字母顺序排序，然后按频率降序排序
        words_and_frequencies = sorted(counter.items(), key=lambda tup: tup[0])
        words_and_frequencies.sort(key=lambda tup: tup[1], reverse=True)

        # 将满足频率要求且不超过最大大小的单词添加到词汇表中
        for word, freq in words_and_frequencies:
            if freq < min_freq or len(self.itos) == max_size:
                break
            self.itos.append(word)

        # stoi 是 itos 的反向字典
        self.stoi = {tok: i for i, tok in enumerate(self.itos)}
        self.vectors = None
        if vectors is not None:
            self.load_vectors(vectors, unk_init=unk_init, cache=vectors_cache)
        else:
            assert unk_init is None and vectors_cache is None

    def __eq__(self, other):
        if self.freqs != other.freqs:
            return False
        if self.stoi != other.stoi:
            return False
        if self.itos != other.itos:
            return False
        if self.vectors != other.vectors:
            return False
        return True

    def __len__(self):
        return len(self.itos)

    def vocab_rerank(self):
        self.stoi = {word: i for i, word in enumerate(self.itos)}

    def extend(self, v, sort=False):
        words = sorted(v.itos) if sort else v.itos
        for w in words:
            if w not in self.stoi:
                self.itos.append(w)
                self.stoi[w] = len(self.itos) - 1


class Vocab(TorchVocab):
    def __init__(self, counter, max_size=None, min_freq=1):
        """
        初始化函数，用于设置词汇表的特殊索引和调用父类的初始化方法。

        :param counter: 一个计数器对象，用于统计词频信息。
        :param max_size: 词汇表的最大容量，默认为None，表示无限制。
        :param min_freq: 词频的最小值，低于此值的词不会被收入词汇表，默认为1。
        """
        # 设置padding的索引值
        self.pad_index = 0
        # 设置未知词的索引值
        self.unk_index = 1
        # 设置句子结束符的索引值
        self.eos_index = 2
        # 设置句子开始符的索引值
        self.sos_index = 3
        # 设置掩码标记的索引值
        self.mask_index = 4
        # 调用父类的初始化方法，传递特殊标记、最大词汇表大小和最小词频限制
        super().__init__(counter, specials=["<pad>", "<unk>", "<eos>", "<sos>", "<mask>"],
                        max_size=max_size, min_freq=min_freq)

    def to_seq(self, sentece, seq_len, with_eos=False, with_sos=False) -> list:
        pass

    def from_seq(self, seq, join=False, with_pad=False):
        pass

    @staticmethod
    def load_vocab(vocab_path: str) -> 'Vocab':
        with open(vocab_path, "rb") as f:
            return pickle.load(f)

    def save_vocab(self, vocab_path):
        with open(vocab_path, "wb") as f:
            pickle.dump(self, f)


# Building Vocab with text files
class WordVocab(Vocab):
    """
    继承自Vocab的WordVocab类，用于构建和处理词汇表。
    
    参数:
    - texts: 用于构建词汇表的文本数据，可以是字符串列表或单个字符串。
    - max_size: 词汇表的最大大小，默认为None，表示无限制。
    - min_freq: 词汇表中单词的最小出现频率，默认为1。
    """
    """
        example:
        with open(args.corpus_path, "r", encoding=args.encoding) as f:
            vocab = WordVocab(f, max_size=args.vocab_size, min_freq=args.min_freq)
    """
    def __init__(self, texts, max_size=None, min_freq=1):
        # 打印构建词汇表的提示信息
        print("Building Vocab")
        # 初始化计数器，用于统计单词出现频率
        counter = Counter()
        # 遍历文本数据，统计单词频率
        for line in tqdm.tqdm(texts):
            # 判断line是否为列表，如果是，则直接赋值给words
            if isinstance(line, list):
                words = line
            else:
                # 如果line不是列表，则进行字符串处理并分割成单词列表
                words = line.replace("\n", "").replace("\t", "").split()
            # 遍历单词列表，更新计数器
            for word in words:
                counter[word] += 1
        # 调用父类Vocab的构造方法，初始化词汇表
        super().__init__(counter, max_size=max_size, min_freq=min_freq)

    def to_seq(self, sentence, seq_len=None, with_eos=False, with_sos=False, with_len=False):
        """
        将句子转换为序列。
        
        参数:
        - sentence: 需要转换的句子，可以是字符串或单词列表。
        - seq_len: 序列长度，默认为None，表示不进行填充或截断。
        - with_eos: 是否在序列末尾添加结束符号，默认为False。
        - with_sos: 是否在序列开头添加开始符号，默认为False。
        - with_len: 是否返回原始序列长度，默认为False。
        
        返回:
        - 序列或(序列, 原始序列长度)的元组，取决于with_len参数的值。
        """
        # 如果sentence是字符串，则将其分割成单词列表
        if isinstance(sentence, str):
            sentence = sentence.split()
        # 将单词转换为对应的索引序列，未知单词使用unk_index
        seq = [self.stoi.get(word, self.unk_index) for word in sentence]
        # 如果with_eos为True，则在序列末尾添加结束符号
        if with_eos:
            seq += [self.eos_index]  # this would be index 1
        # 如果with_sos为True，则在序列开头添加开始符号
        if with_sos:
            seq = [self.sos_index] + seq
        # 记录原始序列长度
        origin_seq_len = len(seq)
        # 根据seq_len参数进行序列的填充或截断
        if seq_len is None:
            pass
        elif len(seq) <= seq_len:
            seq += [self.pad_index for _ in range(seq_len - len(seq))]
        else:
            seq = seq[:seq_len]
        # 根据with_len参数决定返回值
        return (seq, origin_seq_len) if with_len else seq

    def from_seq(self, seq, join=False, with_pad=False):
        """
        将序列转换回句子。
        
        参数:
        - seq: 需要转换的序列，为索引列表。
        - join: 是否将结果连接为字符串，默认为False，返回列表。
        - with_pad: 是否保留填充符号，默认为False，不保留。
        
        返回:
        - 转换后的单词列表或字符串，取决于join参数的值。
        """
        # 将索引转换为对应的单词，未知索引使用"<%d>"格式化
        words = [self.itos[idx]
                 if idx < len(self.itos)
                 else "<%d>" % idx
                 for idx in seq
                 if not with_pad or idx != self.pad_index]
        # 根据join参数决定返回值
        return " ".join(words) if join else words

    @staticmethod
    def load_vocab(vocab_path: str) -> 'WordVocab':
        """
        从文件中加载词汇表。
        
        参数:
        - vocab_path: 词汇表文件的路径。
        
        返回:
        - 加载的WordVocab实例。
        """
        # 打开词汇表文件，并使用pickle.load方法加载WordVocab实例
        with open(vocab_path, "rb") as f:
            return pickle.load(f)


def build():
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("-c", "--corpus_path", required=True, type=str)
    parser.add_argument("-o", "--output_path", required=True, type=str)
    parser.add_argument("-s", "--vocab_size", type=int, default=None)
    parser.add_argument("-e", "--encoding", type=str, default="utf-8")
    parser.add_argument("-m", "--min_freq", type=int, default=1)
    args = parser.parse_args()

    with open(args.corpus_path, "r", encoding=args.encoding) as f:
        vocab = WordVocab(f, max_size=args.vocab_size, min_freq=args.min_freq)

    print("VOCAB SIZE:", len(vocab))
    vocab.save_vocab(args.output_path)
def load_small():
    from dataset import BERTDataset
    with open("data/vocab.small", "rb") as f:
        vocab=pickle.load(f)
    train_dataset = BERTDataset('./data/corpus.small', vocab, seq_len=20,
                                corpus_lines=None, on_memory=True)    
    train_dataset.__getitem__(0)
    print("Train dataset size:", len(train_dataset))   
    '''
        第一行的内容 ['Cuba to Get Rid of Dollars After a Decade', " HAVANA (Reuters) Cubans rushed to change dollars into  local pesos on Tuesday as President Fidel Castro's communist  government prepared to pull the U.S. currency from circulation  more than a decade after it was legalized here."]
        第一行的数量: 2
    ''' 
if __name__ == "__main__":
    build()
    # load_small()